#!/bin/bash

# Usage: ./wafer.sh args.json

set -e

OUT_DIR=${OUT_DIR:-"build"}

WAFER_COMPILER_BASE_DIR=${WAFER_COMPILER_BASE_DIR:-"../.."}
WAFER_COMPILER_BUILD_DIR=${WAFER_COMPILER_BUILD_DIR:-"$WAFER_COMPILER_BASE_DIR/build"}

WAFER_INC_DIR="$WAFER_COMPILER_BASE_DIR/include"
WAFER_LIB_DIR="$WAFER_COMPILER_BASE_DIR/build/lib"

LLVM_PRJ_DIR=${LLVM_PRJ_DIR:-"$WAFER_COMPILER_BASE_DIR/3rdparty/llvm"}

LLVM_CLANG_INC_DIR="$LLVM_PRJ_DIR/build/lib/clang/17/include"

CC=${CC:-"/usr/bin/gcc"}
CXX=${CXX:-"/usr/bin/g++"}

WAFER_OUT_DIR=${WAFER_OUT_DIR:-"$WAFER_COMPILER_BASE_DIR/build"}
WAFER_FRONTEND="$WAFER_OUT_DIR/bin/wafer-frontend"
WAFER_OPT="$WAFER_OUT_DIR/bin/wafer-opt"
CLANGXX="$LLVM_PRJ_DIR/build/bin/clang++"
LLC="$LLVM_PRJ_DIR/build/bin/llc"

LLVM_LIB_DIR=${LLVM_LIB_DIR:-"$LLVM_PRJ_DIR/build/lib"}

CUDA_BASE_DIR=${CUDA_BASE_DIR:-"/usr/local/cuda"}
CUDA_INC_DIR="$CUDA_BASE_DIR/targets/x86_64-linux/include"
CUDA_LIB_DIR="$CUDA_BASE_DIR/targets/x86_64-linux/lib"
CUDAFEXX="$CUDA_BASE_DIR/bin/cudafe++"
CICC="$CUDA_BASE_DIR/nvvm/bin/cicc"
PTXAS="$CUDA_BASE_DIR/bin/ptxas"
FATBINARY="$CUDA_BASE_DIR/bin/fatbinary"
NVLINK="$CUDA_BASE_DIR/bin/nvlink"
CUDA_CRT_STUB="$CUDA_BASE_DIR/bin/crt/link.stub"

OPENSSL_DIR=${OPENSSL_DIR:-"$WAFER_COMPILER_BASE_DIR/3rdparty/openssl"}
OPENSSL_LIB_DIR="$OPENSSL_DIR/lib64"
OPENSSL_INC_DIR="$OPENSSL_DIR/include"

compile_cuda() {
# nvcc using cuda 10.1 version
  if [[ -r "$OUT_DIR/kernel.ptx" ]]; then
    $CC -D__CUDA_ARCH__=300 -E -x c++ \
      -DCUDA_DOUBLE_MATH_FUNCTIONS -D__CUDACC__ -D__NVCC__ \
      "-I$CUDA_INC_DIR" -D__CUDACC_VER_MAJOR__=10 -D__CUDACC_VER_MINOR__=1 -D__CUDACC_VER_BUILD__=243 -include "cuda_runtime.h" \
      -m64 "./main.cu" > "$OUT_DIR/main.cpp1.ii"

    $CICC --c++14 --gnu_version=70500 --allow_managed \
      -arch compute_30 -m64 -ftz=0 -prec_div=1 -prec_sqrt=1 -fmad=1 --include_file_name "main.fatbin.c" -tused -nvvmir-library \
      "$CUDA_BASE_DIR/nvvm/libdevice/libdevice.10.bc" --gen_module_id_file --module_id_file_name "$OUT_DIR/main.module_id" --orig_src_file_name "./main.cu" \
      --gen_c_file_name "$OUT_DIR/main.cudafe1.c" --stub_file_name "$OUT_DIR/main.cudafe1.stub.c" \
      --gen_device_file_name "$OUT_DIR/main.cudafe1.gpu" "$OUT_DIR/main.cpp1.ii" -o "$OUT_DIR/main.ptx"

    cp $OUT_DIR/kernel.ptx $OUT_DIR/main.ptx

    $FATBINARY --create="$OUT_DIR/main.fatbin" -64 "--image=profile=compute_30,file=$OUT_DIR/main.ptx" --embedded-fatbin="$OUT_DIR/main.fatbin.c"

    $CC -E -x c++ -D__CUDACC__ -D__NVCC__  "-I$CUDA_INC_DIR" -D__CUDACC_VER_MAJOR__=10 -D__CUDACC_VER_MINOR__=1 -D__CUDACC_VER_BUILD__=243 \
      -include "cuda_runtime.h" -m64 "./main.cu" > "$OUT_DIR/main.cpp4.ii"

    $CUDAFEXX --c++14 --gnu_version=70500 --allow_managed  --m64 --parse_templates \
      --gen_c_file_name "$OUT_DIR/main.cudafe1.cpp" --stub_file_name "main.cudafe1.stub.c" \
      --module_id_file_name "$OUT_DIR/main.module_id" "$OUT_DIR/main.cpp4.ii"

    $CC -D__CUDA_ARCH__=300 -c -x c++ -DCUDA_DOUBLE_MATH_FUNCTIONS "-I$CUDA_INC_DIR" -m64 -o "$OUT_DIR/main.o" "$OUT_DIR/main.cudafe1.cpp"

    $NVLINK --arch=compute_30 --register-link-binaries="$OUT_DIR/dlink.reg.c"  -m64 "-L$CUDA_LIB_DIR/stubs" "-L$CUDA_LIB_DIR" -cpu-arch=X86_64 "$OUT_DIR/main.o"

    $FATBINARY --create="$OUT_DIR/dlink.fatbin" -64 -link --embedded-fatbin="$OUT_DIR/dlink.fatbin.c"

    $CC -c -x c++ -DFATBINFILE="\"dlink.fatbin.c\"" -DREGISTERLINKBINARYFILE="\"dlink.reg.c\"" -I$OUT_DIR \
      -D__NV_EXTRA_INITIALIZATION= -D__NV_EXTRA_FINALIZATION= -D__CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__ \
      "-I$CUDA_INC_DIR" -D__CUDACC_VER_MAJOR__=10 -D__CUDACC_VER_MINOR__=1 -D__CUDACC_VER_BUILD__=243 \
      -m64 -o "$OUT_DIR/dlink.o" "$CUDA_BASE_DIR/bin/crt/link.stub"

    $CXX -m64 -o "$OUT_DIR/wafer_gpu.out" -Wl,--start-group "$OUT_DIR/dlink.o" "$OUT_DIR/main.o" "-L$CUDA_LIB_DIR/stubs" "-L$CUDA_LIB_DIR" \
      -lcudadevrt  -L$CUDA_BASE_DIR/lib64 -lcudart -lrt -lpthread -ldl  -Wl,--end-group
  fi
}

compile_host() {
  local design_json_file="$1"; shift
  local source_file="$1"; shift
  local output_file="$1"; shift

  $CLANGXX -c "$source_file" -I$WAFER_INC_DIR -o "$OUT_DIR/$output_file.hostctl.o"

  $WAFER_FRONTEND -I $WAFER_INC_DIR -I $LLVM_CLANG_INC_DIR \
    --sst --JsonFilePath "$design_json_file" -c "$source_file" --wafer-to-llvmir -o "$OUT_DIR/$output_file.ll"

  $LLC -mtriple=x86_64-unknown-linux-gnu "$OUT_DIR/$output_file.ll" --filetype=obj -o "$OUT_DIR/$output_file.kernel.o"

  $WAFER_FRONTEND -I $WAFER_INC_DIR -I $LLVM_CLANG_INC_DIR \
    --ptx --JsonFilePath "$design_json_file" -c "$source_file" -o "$OUT_DIR/$output_file.ptx.mlir"

  $WAFER_OPT "$OUT_DIR/$output_file.ptx.mlir" -device-schedule -lower-hyper-memcpy -lower-hyper | \
    $WAFER_OPT -lower-crypto-padmessage | \
    $WAFER_OPT -lower-cryptosha1-without-SHA1ISA | \
    $WAFER_OPT -convert-async-to-llvm | \
    $WAFER_OPT -convert-scf-to-cf | \
    $WAFER_OPT -gpu-kernel-outlining | \
    $WAFER_OPT -convert-arith-to-llvm | \
    $WAFER_OPT -finalize-memref-to-llvm | \
    $WAFER_OPT -convert-vector-to-llvm | \
    $WAFER_OPT -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin{chip=sm_30 features=+ptx32}))' > /dev/null
  if [[ -f "kernel.ptx" ]]; then
    mv kernel.ptx $OUT_DIR/
  fi
}

linkage() {
  local output_file="$1"; shift

  $CLANGXX $OUT_DIR/"$output_file.hostctl.o" $OUT_DIR/"$output_file.kernel.o" \
    -L $LLVM_LIB_DIR -lmlir_c_runner_utils \
    -L $WAFER_LIB_DIR -lwafer_sst_runtime \
    -Wl,-rpath,$LLVM_LIB_DIR:$WAFER_LIB_DIR \
    -o $OUT_DIR/"$output_file"
}

compile_openssl() {
  local source_file="$1"; shift
  local output_file="$1"; shift
  $CXX -DWAFER_OPENSSL "$source_file" -I $OPENSSL_INC_DIR -L $OPENSSL_LIB_DIR -l crypto -o "$OUT_DIR/$output_file"
}

clean() {
  echo "Removing $OUT_DIR"
  rm -rf $OUT_DIR
}

canonicalize_path() {
  local orig_path="$1"; shift
  local rel_path="$1"; shift

  if [[ "$orig_path" == /* ]]; then
    printf "$orig_path"
  else
    printf "$rel_path/$(basename $orig_path)"
  fi
}

main() {
  local args_json="$1"; shift

  if [[ -z "$args_json" ]]; then
    echo "Empty argument" >&2
    exit 1
  fi

  if [[ "$args_json" = "clean" ]]; then
    clean
    exit 0
  fi

  local args_json_parent_path=$(dirname $(realpath "$args_json"))

local source_file="./AESCrack.cpp"
local json_file="./design.json"
local output_file="kt_run_exe"
local compile_type="openssl"

  unset _get_key

  mkdir -p "$OUT_DIR"

  case "$compile_type" in
    openssl)
      compile_openssl "$source_file" "$output_file"
      ;;
    crypto)
      compile_host "$json_file" "$source_file" "$output_file"
      compile_cuda
      linkage "$output_file"
      ;;
    *)
      echo "Invalid compile_type $compile_type" >&2
      ;;
  esac

  echo "Compile successfully, file saved under $OUT_DIR/$output_file"
}

main $@
